import asyncio
import asyncpg
from async_cpu import async_timed
from concurrent.futures.process import ProcessPoolExecutor

product_query = """
    select * from product as p 
    join sku as s on s.product_id = p.product_id
    join product_color as pc on pc.product_color_id = s.product_color_id
    join product_size as ps on ps.product_size_id = s.product_size_id
    where p.product_id = 100
"""


async def query_product(pool):
    async with pool.acquire() as connection:
        return await connection.fetchrow(product_query)


@async_timed()
async def query_products_concurrnetly(pool, queries):
    queries = [query_product(pool) for _ in range(queries)]
    return await asyncio.gather(*queries)


def run_in_new_loop(num_queries):
    async def run_requires():
        async with asyncpg.create_pool(
                host="127.0.0.1",
                port=5432,
                user="around",
                password="zhouwei",
                database="products",
                min_size=6,
                max_size=6,
        ) as pool:
            return await query_products_concurrnetly(pool, num_queries)

    # 在一个新的事件循环中进行查询，并将其转换为字典
    results = [dict(result) for result in asyncio.run(run_requires())]
    return results


@async_timed()
async def main():
    loop = asyncio.get_running_loop()
    pool = ProcessPoolExecutor()
    # 创建5个进程，每个进程都有自己的事件循环来运行查询
    tasks = [loop.run_in_executor(pool, run_in_new_loop, 10000) for _ in range(5)]
    # 等待所有查询结果完成
    all_results = await asyncio.gather(*tasks)

    total_queries = sum([len(result) for result in all_results])
    print(total_queries)


if __name__ == '__main__':
    asyncio.run(main())
